from datetime import datetime, time as datetime_time, timedelta
import time
import json
import requests
from selenium import webdriver
from selenium.webdriver.common.by import By

STATION_DATA_FILE = "stations-g.json"
CITY_LIST_FILE = "provincial_capitals-s.txt"
DEFAULT_DATE = time.strftime("%Y-%m-%d", time.localtime(time.time() + 3 * 24 * 60 * 60))
RETRY_DELAY = 3

# Format of detailed data of every city:
# Train number, Departure city, Arrival city, Departure station, Arrival station, Departure time, Arrival time, Duration, Op on Day 1, Op on Day 2, Op on Day 3, Op on Day 4, Op on Day 5, Op on Day 6, Op on Day 7

# Define the data structure for the train information
class TrainInfo:
	def __init__(self, train_no, dept_city, arrv_city, dept_station, arrv_station, dept_time, arrv_time, duration, op_days):
		self.train_no = train_no			# Type: string
		self.dept_city = dept_city			# Type: string (of city name)
		self.arrv_city = arrv_city
		self.dept_station = dept_station	# Type: string (of station code)
		self.arrv_station = arrv_station
		self.dept_time = dept_time			# Type: datetime_time
		self.arrv_time = arrv_time			# Type: datetime_time
		self.duration = duration			# Type: datetime_time
		self.op_days = op_days				# Type: list[bool]

# Get the train information from `dcity` to `acity` on `date` + `date_offset`
# Return a list of `TrainInfo` objects and a boolean indicating whether the query is successful
# `dcity` and `acity` are city names
# `date` is a string in the format "YYYY-MM-DD" and `date_offset` is an integer from 0 to 6
def get_train_data_for_one_day(dcity, acity, date, date_offset, driver, station_data):
	assert dcity in station_data, f"Invalid departure city: {dcity}"
	assert acity in station_data, f"Invalid arrival city: {acity}"
	assert date >= datetime.now().strftime("%Y-%m-%d"), f"Invalid start date: {date}"
	assert driver is not None, "Driver is not initialized"

	# Add the date offset
	date = datetime.strptime(date, "%Y-%m-%d") + timedelta(days=date_offset)
	date = date.strftime("%Y-%m-%d")
	# Result of the query in string format
	qresult = []
	# Train list of `TrainInfo` objects
	train_list = []
	# Concatenate the URL
	url = f'https://kyfw.12306.cn/otn/leftTicket/queryG?leftTicketDTO.train_date={date}&leftTicketDTO.from_station={station_data[dcity]}&leftTicketDTO.to_station={station_data[acity]}&purpose_codes=ADULT'
	try:
		driver.get(url)
		body = driver.find_elements(By.TAG_NAME, "body")[0].get_attribute('innerHTML')
		# get index of "[" just behind "result"
		start = body.find("[", body.find("\"result\"")) + 1
		if start == 0:
			# No train information found, failed!
			return [train_list, False]
		# get index of "]", which is the end of the JSON data
		end = body.find("]", start)
		# extract the JSON data
		data = body[start:end] # data is a string, without the outermost "[" and "]"
		data = data.split(",")
		for i in range(0, len(data), 2):
			qresult.append(data[i][1:-2])
	except (requests.RequestException, KeyError):
		print("Failed to get train information, please check the network or URL address!")
		exit(1)

	for train_str in qresult:
		try:
			train_info = train_str.split("|")
			# All of them are strings
			train_no = train_info[3]
			dept_station = train_info[6]	# Type: string (of station code)	
			arrv_station = train_info[7]
			dept_time = train_info[8]
			arrv_time = train_info[9]
			duration = train_info[10]
			# Format the departure, arrival time and duration
			dept_time = datetime.strptime(dept_time, "%H:%M").time()
			arrv_time = datetime.strptime(arrv_time, "%H:%M").time()
			# `duration`` is in the format "HH:MM", but it may exceed 24 hours!
			# If it exceeds 24 hours, set it as "00:00"
			try:
				duration = datetime.strptime(duration, "%H:%M").time()
			except ValueError:
				duration = datetime.strptime("00:00", "%H:%M").time()
		# Catch any exceptions and skip this train
		except (ValueError, IndexError):
			# If any of the above fails, skip this train
			continue
		# Set `op_days` to a list of 7 booleans only True at `date_offset`
		op_days = [False] * 7
		op_days[date_offset] = True
		# Append this train to the list
		train_list.append(TrainInfo(train_no, dcity, acity, dept_station, arrv_station, dept_time, arrv_time, duration, op_days))
	return [train_list, True]

# Get the train information from `dcity` to `acity` on the week of `date`
# Return a list of `TrainInfo` objects
# `dcity` and `acity` are city names
# `date` is a string in the format "YYYY-MM-DD"
def get_train_data_for_week(dcity, acity, date, driver, station_data):
	train_list_one_week = []
	for day_offs in range(7):
		flag = False
		# Try until the query is successful
		while not flag:
			[train_list_one_day, flag] = get_train_data_for_one_day(dcity, acity, date, day_offs, driver, station_data)
			if not flag:
				# If failed, wait for 1 second and retry
				print(f"Failed to get train information for {dcity} to {acity} on {date} + {day_offs}, sleeping...")
				time.sleep(RETRY_DELAY)
		for train in train_list_one_day:
			# Search the train number, departure station and arrival station in `train_list_one_week`
			flag_found = False
			for t in train_list_one_week:
				if t.train_no == train.train_no and t.dept_station == train.dept_station and t.arrv_station == train.arrv_station:
					flag_found = True
					t.op_days[day_offs] = True
			if not flag_found:
				# If not found, append a new train
				train_list_one_week.append(train)
	return train_list_one_week

def get_train_data_for_a_source_city(scity, date, driver, station_data):
    # Open the provincial capitals file
	with open(CITY_LIST_FILE, "r", encoding="utf-8") as f:
		prov_caps = f.readlines()
	# Remove the newline characters
	prov_caps = [city.strip() for city in prov_caps]
	train_list = []
	for acity in prov_caps:
		if acity != scity:
			print(f"Getting train information from {scity} to {acity} on {date}+...")
			train_list += get_train_data_for_week(scity, acity, date, driver, station_data)
	return train_list

# Save the train information (in a list of `TrainInfo`) to a CSV file
def save_train_data_to_csv(train_list, output_file, station_data):
	with open(output_file, "w", encoding="utf-8") as f:
		f.write("Train No,Departure City,Arrival City,Departure Station,Arrival Station,Departure Time,Arrival Time,Duration,Op on Day 1,Op on Day 2,Op on Day 3,Op on Day 4,Op on Day 5,Op on Day 6,Op on Day 7\n")
		for train in train_list:
			# Get the station names from the station codes
			try:
				dept_station_name = list(station_data.keys())[list(station_data.values()).index(train.dept_station)]
			except ValueError:
				dept_station_name = train.dept_station
			try:
				arrv_station_name = list(station_data.keys())[list(station_data.values()).index(train.arrv_station)]
			except ValueError:
				arrv_station_name = train.arrv_station
			# Format the time to string and write
			f.write(f"{train.train_no},{train.dept_city},{train.arrv_city},{dept_station_name},{arrv_station_name},{train.dept_time.strftime("%H:%M")},{train.arrv_time.strftime("%H:%M")},{train.duration.strftime("%H:%M")}")
			for op_day in train.op_days:
				# Write 1 if the train operates on that day, 0 otherwise
				f.write(",1" if op_day else ",0")
			f.write("\n")

def init_driver():
	driver = webdriver.Edge()
	driver.get(
		f'https://kyfw.12306.cn/otn/leftTicket/init?linktypeid=dc&fs=%E6%9D%AD%E5%B7%9E,HZH&ts=%E9%93%B6%E5%B7%9D,YIJ&date={DEFAULT_DATE}&flag=N,N,Y')
	return driver

def main():
	driver = init_driver()
	with open(STATION_DATA_FILE, "r", encoding="utf-8") as f:
		station_data = json.load(f)
	start_city = input("source station ([q]uit, [A]ll): ")
	if start_city == "q":
		exit(0)
	start_date = input("Departure date(YYYY-MM-DD): ")
	if start_city == "A":
		# Open the provincial capitals file
		with open(CITY_LIST_FILE, "r", encoding="utf-8") as f:
			prov_caps = f.readlines()
		# Remove the newline characters
		prov_caps = [city.strip() for city in prov_caps]
		for city in prov_caps:
			train_list = get_train_data_for_a_source_city(city, start_date, driver, station_data)
			save_train_data_to_csv(train_list, f"Data/{city}-{start_date}-d.csv", station_data)
	else:
		train_list = get_train_data_for_a_source_city(start_city, start_date, driver, station_data)
		save_train_data_to_csv(train_list, f"Data/{start_city}-{start_date}-d.csv", station_data)

if __name__ == '__main__':
	main()
